{
    "cells": [
  {
   "cell_type": "markdown",
   "id": "875e898a",
   "metadata": {},
   "source": [
    "# KISS-ICP running on the Newer College Dataset\n",
    "\n",
    "The whole purpose of this notebook is to have a reproducable entry point for the experiments of the paper. You can also modify the system and just run this notebook to inspect the overall results.\n",
    "\n",
    "We use the \"old\" version of the newer college dataset, the one that uses the OS1-64 LiDAR. There are 2 public sequences, the \"long_experiment\" is only avaiable as rosbag file.\n",
    "\n",
    "## Expected dataset layout\n",
    "\n",
    "```sh\n",
    "newer_college\n",
    "  ├── 01_short_experiment\n",
    "  │   ├── ground_truth\n",
    "  │   │   ├── poses_kitti_format.txt\n",
    "  │   │   └── registered_poses.csv\n",
    "  │   ├── raw_format\n",
    "  │   │   ├── ouster_imu\n",
    "  │   │   ├── ouster_scan\n",
    "  │   │   │   ├── cloud_1583836591_182590976.pcd\n",
    "  │   │   │   ├── cloud_1583836591_282592512.pcd\n",
    "  │   │   │   ├── ...\n",
    "  │   │   └── realsense_imu\n",
    "  │   └── time_offsets\n",
    "  │       └── time_offsets.csv\n",
    "  ├── 02_long_experiment\n",
    "  │   ├── ground_truth\n",
    "  │   │   └── registered_poses.csv\n",
    "  │   ├── raw_format\n",
    "  │   │   └── ouster_scan\n",
    "  │   └── rosbag\n",
    "  │       ├── rooster_2020-03-10-11-36-51_0.bag\n",
    "  │       ├── rooster_2020-03-10-11-39-38_1.bag\n",
    "  │       ├── ...\n",
    "\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16d8e156",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install KISS-ICP and Plotting tools\n",
    "%pip install kiss-icp ipympl evo >/dev/null\n",
    "\n",
    "import os\n",
    "from pathlib import Path\n",
    "\n",
    "import kiss_icp\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from evo.tools import plot\n",
    "from kiss_icp.datasets import dataset_factory\n",
    "from kiss_icp.pipeline import OdometryPipeline\n",
    "\n",
    "%autoreload 2\n",
    "%matplotlib inline\n",
    "%matplotlib widget"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9c1576b",
   "metadata": {},
   "outputs": [],
   "source": [
    "data_root = os.environ.get(\"DATASETS\")\n",
    "data_dir = Path(os.path.join(data_root, \"newer_college/2020-ouster-os1-64-realsense\"))\n",
    "print(f\"Reading datasets from : {data_root}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a9d09d6",
   "metadata": {},
   "source": [
    "## Custom Dataloader for the `02_long_experiment` sequence\n",
    "\n",
    "For whatever reason this part of the dataset is only public as rosbag. So we need to do a bit of extra effort to run the experiment (but not too much ;)\n",
    "\n",
    "It is also a nice exercise on how can you extend KISS-ICP to process data from your own dataset by simply providing a custom dataloader (just one python class!)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "359ebd2a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import rosbag\n",
    "import sensor_msgs.point_cloud2 as pc2\n",
    "from natsort import natsorted\n",
    "from pyquaternion import Quaternion\n",
    "\n",
    "\n",
    "class NewerCollegeRosbag:\n",
    "    def __init__(self, data_dir: Path):\n",
    "        self.data_source = os.path.join(data_dir, \"\")\n",
    "        self.pose_file = os.path.join(\n",
    "            self.data_source, \"ground_truth/registered_poses.csv\"\n",
    "        )\n",
    "        self.gt_poses = self.load_gt_poses(self.pose_file)\n",
    "        self.sequence_id = \"02_long_experiment\"\n",
    "\n",
    "        # process rosbag-files\n",
    "        print(\"Processing rosbag files (this will take some time)\")\n",
    "        self.topic = \"/os1_cloud_node/points\"\n",
    "        self.rosbag_dir = data_dir / \"rosbag\"\n",
    "        self.rosbag_files = natsorted(\n",
    "            [\n",
    "                os.path.join(self.rosbag_dir, f)\n",
    "                for f in os.listdir(self.rosbag_dir)\n",
    "                if f.endswith(\".bag\")\n",
    "            ]\n",
    "        )\n",
    "        self.bags = []\n",
    "        self.n_scans = 0\n",
    "        for rosbag_file in self.rosbag_files:\n",
    "            bag = rosbag.Bag(rosbag_file, mode=\"r\")\n",
    "            self.bags.append(bag)\n",
    "            self.n_scans += bag.get_message_count(topic_filters=self.topic)\n",
    "\n",
    "        self.bag = self.bags.pop(0)\n",
    "        self.msgs = self.bag.read_messages(topics=[self.topic])\n",
    "        print(f\"Now processing {self.bag.filename}\")\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.n_scans\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        try:\n",
    "            _, msg, _ = next(self.msgs)\n",
    "        except StopIteration:\n",
    "            self.bag.close()\n",
    "            # new bagile\n",
    "            self.bag = self.bags.pop(0)\n",
    "            self.msgs = self.bag.read_messages(topics=[self.topic])\n",
    "            print(f\"Now processing {self.bag.filename}\")\n",
    "            _, msg, _ = next(self.msgs)\n",
    "\n",
    "        points = np.array(list(pc2.read_points(msg, field_names=[\"x\", \"y\", \"z\"])))\n",
    "        return points.astype(np.float64)\n",
    "\n",
    "    @staticmethod\n",
    "    def load_gt_poses(file_path: str):\n",
    "        \"\"\"Taken from pyLiDAR-SLAM/blob/master/slam/dataset/nhcd_dataset.py\"\"\"\n",
    "        ground_truth_df = np.genfromtxt(str(file_path), delimiter=\",\", dtype=np.float64)\n",
    "        xyz = ground_truth_df[:, 2:5]\n",
    "        rotations = np.array(\n",
    "            [\n",
    "                Quaternion(x=x, y=y, z=z, w=w).rotation_matrix\n",
    "                for x, y, z, w in ground_truth_df[:, 5:]\n",
    "            ]\n",
    "        )\n",
    "\n",
    "        num_poses = rotations.shape[0]\n",
    "        poses = np.eye(4, dtype=np.float64).reshape(1, 4, 4).repeat(num_poses, axis=0)\n",
    "        poses[:, :3, :3] = rotations\n",
    "        poses[:, :3, 3] = xyz\n",
    "\n",
    "        T_CL = np.eye(4, dtype=np.float32)\n",
    "        T_CL[:3, :3] = Quaternion(x=0, y=0, z=0.924, w=0.383).rotation_matrix\n",
    "        T_CL[:3, 3] = np.array([-0.084, -0.025, 0.050], dtype=np.float32)\n",
    "        poses = np.einsum(\"nij,jk->nik\", poses, T_CL)\n",
    "        poses = np.einsum(\"ij,njk->nik\", np.linalg.inv(poses[0]), poses)\n",
    "        return poses"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5be26677",
   "metadata": {},
   "outputs": [],
   "source": [
    "from kiss_icp_eval import run_sequence\n",
    "\n",
    "\n",
    "def ncd_short_sequence(int):\n",
    "    return OdometryPipeline(\n",
    "        dataset=dataset_factory(\n",
    "            dataloader=\"ncd\",\n",
    "            data_dir=data_dir / \"01_short_experiment\",\n",
    "        ),\n",
    "    )\n",
    "\n",
    "\n",
    "def ncd_long_sequence(int):\n",
    "    return OdometryPipeline(\n",
    "        dataset=NewerCollegeRosbag(data_dir=data_dir / \"02_long_experiment\"),\n",
    "    )\n",
    "\n",
    "\n",
    "results = {}\n",
    "run_sequence(ncd_short_sequence, sequence=\"01_short_experiment\", results=results)\n",
    "run_sequence(ncd_long_sequence, sequence=\"02_long_experiment\", results=results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a07bc96f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from kiss_icp_eval import print_metrics_table\n",
    "\n",
    "print_metrics_table(results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "03607ea4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from kiss_icp_eval import plot_trajectories\n",
    "\n",
    "plot_trajectories(results)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.6"
  },
  "vscode": {
   "interpreter": {
    "hash": "31f2aee4e71d21fbe5cf8b01ff0e069b9275f58929596ceb00d14d90e3e16cd6"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
